"""
weakly_main.py - 弱监督音视频事件检测训练主程序

实现基于弱监督学习的音视频事件检测模型训练流程，包含数据加载、模型训练、验证评估等功能。
"""

import os
import time
import random
import json
from tqdm import tqdm

import torch
import torch.nn as nn
from torch.nn.utils import clip_grad_norm_
from torch.utils.data import DataLoader
from tensorboardX import SummaryWriter
from torch.optim.lr_scheduler import StepLR, MultiStepLR

import numpy as np
from configs.opts import parser
from model.main_model import weak_main_model as main_model
from utils import AverageMeter, Prepare_logger, get_and_save_args
from utils.Recorder import Recorder
from dataset.AVE_dataset_weak import AVEDataset

# ================================= 随机种子配置 ============================
SEED = 666
random.seed(SEED)
np.random.seed(seed=SEED)
torch.manual_seed(seed=SEED)
torch.cuda.manual_seed(seed=SEED)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

# =============================================================================
# 加载模型配置文件
config_path = 'configs/weak.json'
with open(config_path) as fp:
    config = json.load(fp)
print(config)


def main():
    """
    主函数：程序执行入口
    
    功能：
    - 初始化配置参数和日志系统
    - 加载训练和测试数据集
    - 初始化模型、优化器和损失函数
    - 执行训练和验证循环
    - 保存最佳模型检查点
    """
    # 全局变量声明
    global args, logger, writer, dataset_configs
    global best_accuracy, best_accuracy_epoch
    
    # 初始化最佳准确率记录变量
    best_accuracy, best_accuracy_epoch = 0, 0
    
    # 解析命令行参数和配置文件
    dataset_configs = get_and_save_args(parser)
    parser.set_defaults(**dataset_configs)
    args = parser.parse_args()
    
    # 设置GPU环境变量
    os.environ['CUDA_DEVICE_ORDER'] = "PCI_BUS_ID"
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

    # 创建模型保存目录
    if not os.path.exists(args.snapshot_pref):
        os.makedirs(args.snapshot_pref)

    # 如果从检查点恢复，更新保存目录路径
    if os.path.isfile(args.resume):
        args.snapshot_pref = os.path.dirname(args.resume)

    # 初始化日志记录器
    logger = Prepare_logger(args, eval=args.evaluate)

    # 输出运行时信息
    if not args.evaluate:
        logger.info(f'\n创建模型保存目录: {args.snapshot_pref}')
        logger.info('\n运行时参数:\n\n{}\n'.format(json.dumps(vars(args), indent=4)))
    else:
        logger.info(f'\n评估日志将保存至: {args.snapshot_pref}/Eval.log')

    # 初始化训练数据加载器
    train_dataloader = DataLoader(
        AVEDataset('./data/', split='train'),  # 训练集
        batch_size=args.batch_size,
        shuffle=True,  # 训练时打乱数据顺序
        num_workers=8,  # 数据加载进程数
        pin_memory=True  # 锁页内存，加速GPU传输
    )

    # 初始化测试数据加载器
    test_dataloader = DataLoader(
        AVEDataset('./data/', split='test'),  # 测试集
        batch_size=args.test_batch_size,
        shuffle=False,  # 测试时不打乱数据顺序
        num_workers=8,
        pin_memory=True
    )

    # 初始化主模型
    mainModel = main_model(config["model"])
    # 多GPU数据并行
    mainModel = nn.DataParallel(mainModel).cuda()
    # 获取模型可学习参数
    learned_parameters = mainModel.parameters()
    # 初始化Adam优化器
    optimizer = torch.optim.Adam(learned_parameters, lr=args.lr)
    # 初始化多步长学习率调度器
    scheduler = MultiStepLR(optimizer, milestones=[10, 20, 40], gamma=0.5)
    # 初始化二分类损失函数
    criterion = nn.BCEWithLogitsLoss().cuda()
    # 初始化多标签分类损失函数
    criterion_event = nn.MultiLabelSoftMarginLoss().cuda()

    # 从检查点恢复模型权重
    if os.path.isfile(args.resume):
        logger.info(f"\n加载模型检查点: {args.resume}\n")
        mainModel.load_state_dict(torch.load(args.resume))
    elif args.resume != "" and (not os.path.isfile(args.resume)):
        # 检查点文件不存在时抛出异常
        raise FileNotFoundError(f"检查点文件不存在: {args.resume}")

    # 仅执行评估模式
    if args.evaluate:
        logger.info(f"\n开始模型评估...")
        validate_epoch(mainModel, test_dataloader, criterion, criterion_event, epoch=0, eval_only=True)
        return

    # 初始化Tensorboard日志记录器
    writer = SummaryWriter(args.snapshot_pref)
    # 初始化代码备份记录器
    recorder = Recorder(args.snapshot_pref, ignore_folder="Exps/")
    # 保存运行参数配置
    recorder.writeopt(args)

    # 训练循环
    for epoch in range(args.n_epoch):
        # 执行一个训练周期
        loss = train_epoch(mainModel, train_dataloader, criterion, criterion_event, optimizer, epoch)

        # 定期验证或最终周期验证
        if ((epoch + 1) % args.eval_freq == 0) or (epoch == args.n_epoch - 1):
            # 执行验证周期
            acc = validate_epoch(mainModel, test_dataloader, criterion, criterion_event, epoch)
            
            # 更新最佳准确率并保存模型
            if acc > best_accuracy:
                best_accuracy = acc
                best_accuracy_epoch = epoch
                # 保存最佳模型检查点
                save_checkpoint(
                    mainModel.state_dict(),
                    top1=best_accuracy,
                    task='WeaklySupervised',
                    epoch=epoch + 1,
                )
            
            # 输出当前最佳结果
            print("-----------------------------")
            print("最佳准确率及轮次:", best_accuracy, best_accuracy_epoch)
            print("-----------------------------")
        
        # 更新学习率
        scheduler.step()


def train_epoch(model, train_dataloader, criterion, criterion_event, optimizer, epoch):
    """
    执行单个训练周期
    
    参数:
        model: 待训练的神经网络模型
        train_dataloader: 训练数据加载器
        criterion: 损失函数
        criterion_event: 事件分类损失函数  
        optimizer: 模型参数优化器
        epoch: 当前训练周期数
        
    返回:
        losses.avg: 当前周期的平均损失值
    """
    # 初始化统计指标
    batch_time = AverageMeter()  # 批次处理时间统计
    data_time = AverageMeter()   # 数据加载时间统计
    losses = AverageMeter()      # 损失值统计
    train_acc = AverageMeter()   # 训练准确率统计
    end_time = time.time()       # 当前时间戳

    # 设置模型为训练模式
    model.train()
    # 设置模型参数为双精度浮点数
    model.double()
    # 清空优化器梯度
    optimizer.zero_grad()

    # 遍历训练数据批次
    for n_iter, batch_data in enumerate(train_dataloader):
        # 更新数据加载时间统计